import matplotlib.pyplot as plt
import pandas as pd
from sklearn.metrics import roc_curve, auc

from src.features.FeatureEngineering import feature_engineering
def auc_plot(y_test,y_pred_proba):
    # 计算FPR、TPR和ROC曲线
    fpr, tpr, thresholds = roc_curve(y_test, y_pred_proba)
    # 计算AUC
    auc_value = auc(fpr, tpr)
    # 绘制ROC曲线
    plt.plot(fpr, tpr, label='ROC curve (AUC = {:.2f})'.format(auc_value))
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title('Receiver Operating Characteristic (ROC) Curve')
    plt.legend()
    plt.show()
def data_processing(code=1):
    '''
    参数
    :param code:1:训练集 0:测试集 default:训练集
    :return:x,y(特征,标签)
    '''
    path = r'../../data/raw/train.csv'
    if code == 0:
        path = r'../../data/raw/test2.csv'
    data_ready = pd.read_csv(path)
    data_ready = pd.get_dummies(data_ready)
    data,feature_columns = feature_engineering(data_ready)
    if code == 0:
        data.to_csv(r'../../data/processed/test_ready.csv', index=False)
    else:
        data.to_csv(r'../../data/processed/train_ready.csv', index=False)
    # print(data_ready.info())
    y = data.Attrition
    x = data[feature_columns]
    return x, y


if __name__ == '__main__':
    x, y = data_processing(1)
    print(x)
    print(y)
